# pymq - The python based Message Queue
#
# Copyright(C) 2010, Dhruv Matani(dhuvbird@gmail.com)
#
# pymq is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of
# the License, or (at your option) any later version.
#
# pymq is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with pymq. If not, see
# <http://www.gnu.org/licenses/>.
#


import MySQLdb
import threading

dbPrefix        = "pymq_"
DB_HOST         = 'localhost'
DB_USERNAME     = 'root'
DB_PASSWORD     = ''
MASTER_DB_NAME  = "master"


queryCreateQueuesTable = """CREATE TABLE IF NOT EXISTS
    %s.queues(messageID INTEGER PRIMARY KEY AUTO_INCREMENT NOT NULL,
    queueName CHAR(40) NOT NULL, message TEXT NOT NULL)
    ENGINE=MyISAM DEFAULT CHARSET=utf8"""

queryCreateQueuesIndex = """CREATE INDEX
    queues_queueName ON %s.queues(queueName, messageID)"""


def setConnectionString(connectionString):
    """
    Set the connection string to use when connecting to the data store
    using the DataStore class class below
    """
    global DB_HOST, DB_USERNAME, DB_PASSWORD
    creds, host = connectionString.split("@")
    username, password = creds.split(":")
    DB_HOST = host
    DB_USERNAME = username
    DB_PASSWORD = password


class DataStore(object):
    def __init__(self, dbName, autocommit=True):
        """
        Initializes the DataStore with a database named dbName
        """
        if not hasattr(self, "db"):
            self.db = None
        if not hasattr(self, "conn"):
            self.conn = None
        self.dbName = dbName
        self.autocommit = autocommit

        try:
            self.db = MySQLdb.connect(host=DB_HOST, user=DB_USERNAME,
                                      passwd=DB_PASSWORD, db=dbPrefix + dbName)
            if self.db:
                self.db.autocommit(autocommit)
                self.conn = MySQLdb.cursors.DictCursor(self.db)
        except Exception, e:
            print "Exception in Connection(): %s" % str(e)

    def __call__(self):
        """
        Returns True if this DataStore is usable and False otherwise
        """
        return self.db and self.conn

    def execute(self, query, params=(), tries=0):
        """
        Executes a MySQL query

        Returns None if the query failed, or the return code of the
        underlying cursor object's execute method
        """
        if tries > 1:
            return None
        try:
            return self.conn.execute(query, params)
        except MySQLdb.OperationalError, oe:
            # Try reconnect here (only in case of an OperationalError)
            fullQuery = query % params
            print "OperationalError while executing query (%s)" % fullQuery
            print "Trying to reconnect...."
            if tries == 0:
                self.__init__(self.dbName, self.autocommit)
                if self():
                    return self.execute(query, params, tries+1)

        except Exception, e:
            fullQuery = query % params
            print "Exception while executing query (%s)" % fullQuery
            return None

        return None

    def fetchall(self):
        try:
            self.conn.fetchall()
            return map(lambda x: x, self.conn)
        except Exception, e:
            print "Exception (%s) fetching result set" % str(e)
        return None

    def allowOnlyMaster(self):
        """
        Throws an exception if this isn't the master datastore
        """
        if self.dbName != MASTER_DB_NAME:
            raise RuntimeError("Only the master datastore can add a new datastore")

    def rowcount(self):
        return self.conn.rowcount

    # The commands below are callable only by the master DB handler (ie. the admin)
    def addDataStore(self, dataStoreId):
        """
        Provision and add a new data store on the underlying storage media.
        This actually adds a new data base to the data store.
        """
        self.allowOnlyMaster()
        q = "CREATE DATABASE IF NOT EXISTS " + (dbPrefix + dataStoreId)
        createDBResult = self.execute(q)
        if createDBResult is None:
            return None
        createTableResult = self.execute(queryCreateQueuesTable % (dbPrefix + dataStoreId))
        if createTableResult is None:
            return None
        createIndexResult = self.execute(queryCreateQueuesIndex % (dbPrefix + dataStoreId))
        if createIndexResult is None:
            return None
        q = "INSERT INTO pymqDatabases VALUES(%s)"
        return self.execute(q, (dataStoreId, ))


    def getAllQueueDataStores(self):
        """
        Fetches names of all Queue Data Stores in the system
        """
        self.allowOnlyMaster()
        q = "SELECT * FROM pymqDatabases"
        self.execute(q)
        return self.fetchall()

    def getAllQueuesMetadata(self):
        """
        Get the metadata (queueName, ownerDB, status) for all the
        queues in the system
        """
        self.allowOnlyMaster()
        self.execute("SELECT queueName, ownerDB, status FROM queueMetadata")
        return self.fetchall()


    def addQueue(self, queueName, targetDBId, queueStatus):
        """
        Add a new queue to the system
        """
        self.allowOnlyMaster()
        q = "INSERT INTO queueMetadata(queueName, ownerDB, status) VALUES(%s, %s, %s)"
        return self.execute(q, (queueName, targetDBId, queueStatus))

    def dropQueue(self, queueName):
        """
        Drop a queue from the system
        """
        self.allowOnlyMaster()
        q = "DELETE FROM queueMetadata WHERE queueName=%s"
        return self.execute(q, (queueName, ))

    def setQueueStatus(self, queueName, newStatus):
        """
        Set the status for a queue in the system.
        Status can be one of:
        0: Running
        1: Reject Consumers
        2: Reject Producers
        3: Reject All
        """
        self.allowOnlyMaster()
        q = "UPDATE queueMetadata SET status=%s WHERE queueName=%s"
        return self.execute(q, (newStatus, queueName))

    def getQueueMetadata(self, queueName):
        """
        Get the metadata for queue 'queueName' in the system
        """
        self.allowOnlyMaster()
        q = "SELECT * FROM queueMetadata WHERE queueName=%s"
        self.execute(q, (queueName,))
        return self.fetchall()

    # The commands below are callable only by non-master DB handlers
    # (ie. the individual queue's IO Units)
    def enqueueMessage(self, queueName, message):
        """
        Enqueue a message in the queue 'queueName'
        """
        q = "INSERT INTO queues(queueName, message) VALUES(%s, %s)"
        return self.execute(q, (queueName, message))

    def lastInsertedMessageID(self):
        """
        Get the ID of the last message added to this data store
        """
        q = "SELECT LAST_INSERT_ID() AS LID"
        # The query below should NOT fail.
        self.execute(q)
        return self.fetchall()[0]['LID']

    def peekMessage(self, queueName):
        """
        Fetch the first(oldest) un-acknowledged message in the queue 'queueName'
        """
        q = "SELECT * FROM queues WHERE queueName=%s ORDER BY messageID ASC LIMIT 1"
        self.execute(q, (queueName, ))
        return self.fetchall()

    def peekMessages(self, queueName, n):
        """
        Fetch 'n' oldest un-acknowledged message in the queue 'queueName'
        """
        q = "SELECT * FROM queues WHERE queueName=%s ORDER BY messageID ASC LIMIT %s"
        self.execute(q, (queueName, n))
        return self.fetchall()

    def ackMessage(self, queueName, messageID):
        """
        Acknowledge the message with ID 'messageID'.
        A message may be acknowledged multiple times.
        The acknowledgement should succeed all the time.
        """
        self.execute("DELETE FROM queues WHERE messageID=%s", (messageID, ))

    def getLastUnAckedMessageID(self, queueName):
        """
        Get the messageID of the last(oldest) un-acknowledged message
        from the queue 'queueName'
        """
        q = "SELECT MIN(messageID) FROM queues WHERE queueName=%s"
        self.execute(q, (queueName, ))
        return self.fetchall()

    def unprovisionQueue(self, queueName):
        """
        Delete all the messages from the queue 'queueName'
        """
        q = "DELETE FROM queues WHERE queueName=%s"
        self.execute(q, (queueName,))
        return True

    def provisionQueue(self, queueName):
        """
        Provision a new queue 'queueName' in the current datastore
        """
        return True

    def loadQueue(self, owner, request, fileName):
        """
        Load a queue from the file 'fileName'
        owner: The owner of the request
        request: The request that generated this command
        fileName: The file to load the queue from. This file will be placed in the same
        location that dumpQueue placed it in
        """
        queue_loader(owner, request, self.dbName, fileName)

    def dumpQueue(self, owner, request, fileName):
        """
        Dump a queue to the file 'fileName'
        owner: The owner of the request
        request: The request that generated this command
        fileName: The file to dump the queue to. This file will be placed in the same
        location that loadQueue should read it from
        """
        queue_dumper(owner, request, self.dbName, fileName)


class queue_loader:
    def __init__(self, owner, request, dbName, fileName):
        self.owner = owner
        self.request = request
        self.dbName = dbName
        self.fileName = fileName
        self.conn = DataStore(dbName)
        threading.Thread(target=self).start()

    def __call__(self):
        # self.request.pymq_request is guaranteed to not cause any SQL injection
        q = """LOAD DATA INFILE %s INTO TABLE queues CHARACTER SET utf8
               FIELDS TERMINATED BY ',' ESCAPED BY '"' ENCLOSED BY '"'
               LINES TERMINATED BY '\\n' (message)
               SET queueName = '""" + self.request.pymq_queue + """'"""
        result = self.conn.execute(q, ("/tmp/" + self.fileName, ))
        # TODO: !!CHECK THIS!!
        if result is None:
            self.request.error_enqueue(470)
        else:
            self.request.pymq_type = "loadQueueDone"
            self.request.pymq_num_affected_rows = self.conn.rowcount()
            self.owner.handleRequest(self.request)


class queue_dumper:
    def __init__(self, owner, request, dbName, fileName):
        self.owner = owner
        self.request = request
        self.dbName = dbName
        self.fileName = fileName
        self.conn = DataStore(dbName)
        threading.Thread(target=self).start()
        # TODO: How to create detached threads

    def __call__(self):
        q = """SELECT message FROM queues WHERE queueName=%s ORDER BY messageID ASC
               INTO OUTFILE %s CHARACTER SET utf8
               FIELDS TERMINATED BY ',' ESCAPED BY '"' ENCLOSED BY '"'
               LINES TERMINATED BY '\\n'"""
        result = self.conn.execute(q, (self.request.pymq_queue, "/tmp/" + self.fileName, ))
        # TODO: !!CHECK THIS!!
        if result is None:
            self.request.error_enqueue(409)
        else:
            self.request.pymq_type = "dumpQueueDone"
            self.request.pymq_num_affected_rows = self.conn.rowcount()
            self.owner.handleRequest(self.request)
